from pyvirtualdisplay import Display

from selenium.webdriver import Firefox
from selenium.webdriver import FirefoxProfile

class FirefoxRequest(Firefox):
    def __init__(self, firefox_profile = None, firefox_binary = None, timeout = 30):
        super(FirefoxRequest, self).__init__(
            firefox_profile = firefox_profile,
            firefox_binary = firefox_binary,
            timeout = timeout
        )

    def request(self, link = None, filename = None):
        if link:
            self.get(link)

            return {'screenshot': self.save_screenshot(filename)}
        else:
            return None
